import os
import random
import time
from typing import List, Tuple
import concurrent.futures  # noqa: F401
from threading import Semaphore  # noqa: F401

import adata as ad
import akshare as ak
import baostock as bs
import pandas as pd
import pendulum
from stockstats import wrap  # noqa: F401

from trader.tools import market
from trader.tools import timex  # noqa: F401


def excute_batch_fetch_stock(
    stock_codes: List[str],
    start_date: str = "2005-01-01",
    end_date: str = "2050-01-01",
    interval: str = market.K_DAY,
    mode: str = "chunk",
):
    today_dir = pendulum.now().strftime("%Y%m%d")
    if not os.path.exists(f"fs/{today_dir}"):
        os.makedirs(f"fs/{today_dir}")
    """
    1. Determine the number of threads based on the number of CPU cores
    2. Split the stock codes into chunks based on stocks_per_thread
    """
    num_threads = min(len(stock_codes), os.cpu_count() or 8)
    stocks_per_thread = len(stock_codes) // num_threads + (
        len(stock_codes) % num_threads != 0
    )
    stock_chunks = [  # noqa: F841
        stock_codes[i : i + stocks_per_thread]
        for i in range(0, len(stock_codes), stocks_per_thread)
    ]

    def pool_chunk(stock_chunk):
        for stock_code in stock_chunk:
            fetch_single_stock(stock_code, start_date, end_date, interval)
            time.sleep(1)

    def pool_single(stock_code):
        fetch_single_stock(stock_code, start_date, end_date, interval)
        time.sleep(1)

    if mode == "single":
        bs.login()
        for code in stock_codes:
            fetch_single_stock(code, start_date, end_date, interval)
            time.sleep(random.random())
        return

    # 使用线程池并行处理每个股票分片
    with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor:
        bs.login()
        _ = (
            [executor.submit(pool_chunk, chunk) for chunk in stock_chunks]
            if mode == "chunk"
            else [executor.submit(pool_single, code) for code in stock_codes]
        )


def fetch_single_stock(
    stock_code: str, start_date: str, end_date: str, interval: str
) -> Tuple[str, bool]:
    """
    Fetch stock history for a single stock code with error handling.

    Returns:
        Tuple[str, Union[pd.DataFrame, Exception]]:
            A tuple containing the stock code and
            either a DataFrame of fetched data or an Exception if an error occurred.
    """
    today_dir = pendulum.now().strftime("%Y%m%d")
    try:
        df = fetch_stock_history(stock_code, start_date, end_date, interval)
        if df.empty:
            print(f"🌧️  Error fetching data for {stock_code}: Empty DataFrame")
            return (stock_code, False)

        df.to_csv(f"fs/{today_dir}/{stock_code}-{interval}.csv", index=False)
        print(f"📦 Save {stock_code} successfully")
        return (stock_code, True)
    except Exception as e:
        print(f"🔥 Error fetching data for {stock_code}: {e}")
        return (stock_code, False)


def fetch_stock_history(
    stock_code: str = "002444",
    start_date: str = "1970-01-01",
    end_date: str = "2222-01-01",
    interval: str = market.K_DAY,
) -> pd.DataFrame:
    """
    Fetch stock history for a single stock code.
    """
    k_type = market.K_TYPE[interval]

    if interval in [market.K_5m, market.K_15m, market.K_30m, market.K_60m]:
        return _fetch_bs_history(stock_code, start_date, end_date, interval)
        return _fetch_ak_history(stock_code, start_date, end_date, interval)

    return ad.stock.market.get_market(
        stock_code=stock_code, k_type=k_type, start_date=start_date
    )


def _fetch_ak_history(
    stock_code: str = "002444",
    start_date: str = "1970-01-01",
    end_date: str = "2222-01-01",
    interval: str = market.K_DAY,
) -> pd.DataFrame:
    k_type = market.K_TYPE[interval]

    data = ak.stock_zh_a_hist_min_em(
        symbol=stock_code,
        start_date=start_date,
        end_date=end_date,
        period=k_type,
        adjust="qfq",
    )
    rename_columns = {
        "时间": "trade_time",
        "开盘": "open",
        "收盘": "close",
        "最高": "high",
        "最低": "low",
        "涨跌幅": "change_pct",
        "涨跌额": "change",
        "成交量": "volume",
        "成交额": "amount",
        "振幅": "amplitude",
        "换手率": "turnover_ratio",
    }
    result_df = data.rename(columns=rename_columns)
    result_df["stock_code"] = stock_code
    result_df["trade_time"] = pd.to_datetime(result_df["trade_time"]).dt.strftime(
        "%Y-%m-%d %H:%M:%S"
    )
    result_df["trade_date"] = pd.to_datetime(result_df["trade_time"]).dt.strftime(
        "%Y-%m-%d"
    )
    return result_df


def _fetch_bs_history(
    stock_code: str = "002444",
    start_date: str = "1970-01-01",
    end_date: str = "2222-01-01",
    interval: str = market.K_DAY,
) -> pd.DataFrame:
    k_type = market.K_TYPE[interval]

    rs = bs.query_history_k_data_plus(
        market.bs_code(stock_code),
        "date,time,code,open,high,low,close,volume,amount",
        start_date=start_date,
        frequency=str(k_type),
        adjustflag="2",
    )
    data_list = []
    while (rs.error_code == "0") & rs.next():
        data_list.append(rs.get_row_data())

    data = pd.DataFrame(data_list, columns=rs.fields)
    if data.empty:
        return data
    rename_columns = {
        "code": "stock_code",
        "date": "trade_date",
        "time": "trade_time",
    }
    df = data.rename(columns=rename_columns)
    df["stock_code"] = stock_code
    df["trade_time"] = df["trade_time"].apply(
        lambda x: pd.to_datetime(
            f"{x:17}"[:4]
            + "-"
            + f"{x:17}"[4:6]
            + "-"
            + f"{x:17}"[6:8]
            + " "
            + f"{x:17}"[8:10]
            + ":"
            + f"{x:17}"[10:12]
            + ":"
            + f"{x:17}"[12:14]
        )
    )
    df["trade_date"] = pd.to_datetime(df["trade_date"]).dt.strftime("%Y-%m-%d")

    df[["open", "high", "low", "close", "volume", "amount"]] = df[
        ["open", "high", "low", "close", "volume", "amount"]
    ].astype(float)
    df[["open", "high", "low", "close", "volume", "amount"]] = df[
        ["open", "high", "low", "close", "volume", "amount"]
    ].apply(lambda x: x.round(2))
    return df


def fetch_stock_code() -> pd.DataFrame:
    df = ad.stock.info.all_code()
    df["list_date"] = df["list_date"].fillna("1900-01-01")
    df["stock_code"] = df["stock_code"].astype(str)

    # stock_code starts with 688 or 8 or 9 are not tradable
    filter_cond = ~(
        df["stock_code"].str.startswith("68")
        | df["stock_code"].str.startswith("2")
        | df["stock_code"].str.startswith("8")
        | df["stock_code"].str.startswith("9")
        | df["short_name"].str.contains("ST")
    )

    return df[filter_cond]


def test_fetch_stock_history():
    stock_code = "601127"
    res = fetch_stock_history(
        stock_code=stock_code,
        start_date="2023-01-01",
        end_date="2200-01-01",
        interval=market.K_30m,
    )

    stat = market.wrap_stat(res)
    print(
        stat[
            [
                "trade_time",
                "stock_code",
                "close_20_ma",
                "close_34_ma",
                "close_55_ma",
                "close_610_ma",
            ]
        ]
    )


def test_fetch_stock_code():
    res = fetch_stock_code()

    res.to_csv("fs/stock_code.csv", index=False)


if __name__ == "__main__":
    # test_fetch_stock_code()
    # test_fetch_stock_history()

    bs.login()

    start_date = timex.dt_to_date(pendulum.now().subtract(months=3))
    for code in ["002444", "000001", "600000", "300001"]:
        df = fetch_single_stock(
            stock_code=code,
            start_date=start_date,
            end_date="2200-01-01",
            interval=market.K_30m,
        )
        print(df)
    bs.logout()
